#include "main.h"
#include "usart.h"
#include <stdio.h>
#include <string.h>

#include "OTA.h"

uint8_t recv_over_flag = 1;
uint8_t ota_cmd_recv_buf[16];

__asm void MSR_MSP(uint32_t ulAddr)
{
    MSR MSP, r0 // set Main Stack value
    BX r14
}

typedef void (*Jump_Fun)(void);
void IAP_ExecuteApp(uint32_t App_Addr)
{
    Jump_Fun JumpToApp;

    if (((*(__IO uint32_t *)App_Addr) & 0x2FFE0000) == 0x20000000) //检查栈顶地址是否合法.
    {
        JumpToApp = (Jump_Fun) * (__IO uint32_t *)(App_Addr + 4); //用户代码区第二个字为程序开始地址(复位地址)
        MSR_MSP(*(__IO uint32_t *)App_Addr);                      //初始化APP堆栈指针(用户代码区的第一个字用于存放栈顶地址)
        JumpToApp();                                              //跳转到APP.
    }
}

/* flash operations */

#define FLASH_TIMEOUT_VALUE 50000U
#define __DEBUG_FLASH 0

/**
* @brief erase flash sector
*
* @param[in] sector: This parameter must be a value of @ref FLASH_SECTOR_x
* @return   0: success
            other: fail
*/
int8_t flash_erase_sector(uint32_t sector)
{
    HAL_StatusTypeDef status = HAL_OK;
#if __DEBUG_FLASH
    printf("start erase\r\n");
#endif

    status = HAL_FLASH_Unlock();
    if (HAL_OK != status)
    {
        return -1;
    }

    /* Wait for last operation to be completed */
    status = FLASH_WaitForLastOperation((uint32_t)FLASH_TIMEOUT_VALUE);
    if (HAL_OK != status)
    {
        return -2;
    }

    FLASH_EraseInitTypeDef flash_erase_handle;
    flash_erase_handle.TypeErase = FLASH_TYPEERASE_SECTORS;
    flash_erase_handle.Sector = sector;
    flash_erase_handle.NbSectors = 1;
    flash_erase_handle.VoltageRange = FLASH_VOLTAGE_RANGE_3;

    uint32_t SectorError = 0;
    HAL_FLASHEx_Erase(&flash_erase_handle, &SectorError);

    status = HAL_FLASH_Lock();

    if (HAL_OK != status)
    {
#if __DEBUG_FLASH
        printf("SectorError:0x%x\r\n", SectorError);
#endif
        return -3;
    }
#if __DEBUG_FLASH
    printf("start erase\r\n");
#endif
    return 0;
}

int8_t flash_program_word(uint32_t flash_addr, uint32_t *data, uint32_t word_size)
{
    HAL_StatusTypeDef status = HAL_OK;
    uint32_t i;

    status = HAL_FLASH_Unlock();
    if (HAL_OK != status)
    {
        return -1;
    }

    /* Wait for last operation to be completed */
    status = FLASH_WaitForLastOperation((uint32_t)FLASH_TIMEOUT_VALUE);
    if (HAL_OK != status)
    {
        return -2;
    }

    for (i = 0; i < word_size; i++)
    {
        HAL_FLASH_Program(FLASH_TYPEPROGRAM_WORD, (flash_addr + (4 * i)), data[i]);
    }

    status = HAL_FLASH_Lock();

    if (HAL_OK != status)
    {
        return -3;
    }

    return 0;
}

int8_t flash_program_byte(uint32_t flash_addr, const unsigned char *data, uint32_t byte_size)
{
    HAL_StatusTypeDef status = HAL_OK;
    uint32_t i;

    status = HAL_FLASH_Unlock();
    if (HAL_OK != status)
    {
        return -1;
    }

    /* Wait for last operation to be completed */
    status = FLASH_WaitForLastOperation((uint32_t)FLASH_TIMEOUT_VALUE);
    if (HAL_OK != status)
    {
        return -2;
    }

    for (i = 0; i < byte_size; i++)
    {
        HAL_FLASH_Program(FLASH_TYPEPROGRAM_BYTE, (flash_addr + i), data[i]);
    }

    status = HAL_FLASH_Lock();

    if (HAL_OK != status)
    {
        return -3;
    }

    return 0;
}

uint32_t flash_read_word(uint32_t flash_addr, uint32_t *data, uint32_t word_size)
{
    uint32_t i = 0;

    for (i = 0; i < word_size; i++)
    {
        data[i] = *(__IO uint32_t *)(flash_addr + (4 * i));
    }

    return i;
}

uint32_t flash_read_byte(uint32_t flash_addr, uint8_t *data, uint32_t byte_size)
{
    uint32_t i = 0;

    for (i = 0; i < byte_size; i++)
    {
        data[i] = *(__IO uint8_t *)(flash_addr + i);
    }

    return i;
}

/* ota */
void set_bin_size(uint32_t bin_size)
{
    __disable_irq();
    HAL_StatusTypeDef status = HAL_OK;

    status = HAL_FLASH_Unlock();
    if (HAL_OK != status) {
        __enable_irq();
        return;
    }

    HAL_FLASH_Program(FLASH_TYPEPROGRAM_WORD, FLASH_BIN_SIZE_ADDR, bin_size);

    status = HAL_FLASH_Lock();
    
    __enable_irq();
    return;
}

uint32_t get_bin_size(void)
{
    uint32_t bin_size = 0;

    bin_size = *(__IO uint32_t *)(FLASH_BIN_SIZE_ADDR);
    if (bin_size == 0xFFFFFFFF) {
        return 0;
    }

    return bin_size;
}

int8_t erase_app2(void)
{
    __disable_irq();
    flash_erase_sector(FLASH_SECTOR_6);
    flash_erase_sector(FLASH_SECTOR_7);
    __enable_irq();
}

void set_ota_start_flag(void)
{
    __disable_irq();
    HAL_StatusTypeDef status = HAL_OK;

    status = HAL_FLASH_Unlock();
    if (HAL_OK != status) {
        __enable_irq();
        return;
    }

    HAL_FLASH_Program(FLASH_TYPEPROGRAM_BYTE, FLASH_OTA_FLAG_ADDR, OTA_START_FLAG);

    status = HAL_FLASH_Lock();
    __enable_irq();
}

void set_ota_over_flag(void)
{
    __disable_irq();
    HAL_StatusTypeDef status = HAL_OK;

    status = HAL_FLASH_Unlock();
    if (HAL_OK != status) {
        __enable_irq();
        return;
    }

    HAL_FLASH_Program(FLASH_TYPEPROGRAM_BYTE, FLASH_OTA_FLAG_ADDR + 1, OTA_OVER_FLAG);

    status = HAL_FLASH_Lock();
    __enable_irq();
}

static void erase_app1(void)
{
    __disable_irq();
    flash_erase_sector(FLASH_SECTOR_2);
    flash_erase_sector(FLASH_SECTOR_3);
    flash_erase_sector(FLASH_SECTOR_4);
    flash_erase_sector(FLASH_SECTOR_5);
    __enable_irq();
}

#define MOVE_BUF_SIZE 1024

static void move_code(const uint32_t bin_total_size)
{
    uint8_t move_buf[MOVE_BUF_SIZE];
    uint32_t remain_size = bin_total_size;
    uint32_t read_byte_size, read_offset = 0;

    __disable_irq();

    if (MOVE_BUF_SIZE < bin_total_size) {
        read_byte_size = MOVE_BUF_SIZE;
    } else {
        read_byte_size = bin_total_size;
    }

    while (remain_size) {
        flash_read_byte(FLASH_APP_2_ADDR+read_offset, move_buf, read_byte_size);
        flash_program_byte(FLASH_APP_1_ADDR+read_offset, move_buf, read_byte_size);
        read_offset += read_byte_size;

        remain_size -= read_byte_size;
        if (remain_size < MOVE_BUF_SIZE) {
            read_byte_size = remain_size;
        } else {
            read_byte_size = MOVE_BUF_SIZE;
        }
    }
    __enable_irq();
}

void ota_start(void)
{
    int16_t op_ret;
    uint8_t ota_start_flag = 0, oat_over_flag = 0;
    uint32_t ota_bin_size = 0;// uint: byte
    flash_read_byte(FLASH_OTA_FLAG_ADDR, &ota_start_flag, 1);
    flash_read_byte(FLASH_OTA_FLAG_ADDR+1, &oat_over_flag, 1);

    if (OTA_START_FLAG==ota_start_flag && OTA_OVER_FLAG!=oat_over_flag) {
        printf(">>> Ready OAT\r\n");

        /* get ota bin size */
        ota_bin_size = get_bin_size();
        printf("*ota_bin_size: %d\r\n", ota_bin_size);
        if (ota_bin_size == 0) {
            printf(">>> OTA bin size is 0, no OTA\r\n");
            printf(">>> Exit bootload, execute app\r\n");
            IAP_ExecuteApp(FLASH_APP_1_ADDR);
            return;
        }

        /* erase app1 */
        printf(">>> erase app1 start\r\n");
        erase_app1();
        printf(">>> erase app1 end\r\n");

        /* move code */
        printf(">>> move code start\r\n");
        move_code(ota_bin_size);

        set_ota_over_flag();
        printf(">>> OTA success!!!\r\n");
    } else {
        printf(">>> Don't need to OAT\r\n");
    }
    printf(">>> Exit bootload, execute app\r\n");
    IAP_ExecuteApp(FLASH_APP_1_ADDR);
}
